from typing import Dict, Any
from core.context.ctxvars import set_var, get_var, reset_var

CONTEXT_ID = "_context_id"


class Context:
    def __init__(self):
        self.vars: Dict[str, Any] = {}

    def get(self, key: str, default: Any = None):
        return self.vars.get(key, default)

    def set(self, key: str, value: Any):
        self.vars[key] = value

    def clear(self):
        # Clear the context variables
        self.vars.clear()


class ContextManager:
    def set_context(self, ctx: Context):
        set_var(CONTEXT_ID, ctx)

    def get_context(self) -> Context:
        ctx = get_var(CONTEXT_ID)
        if ctx is None:
            ctx = Context()
            set_var(CONTEXT_ID, ctx)
        return ctx

    def release_context(self):
        ctx = get_var(CONTEXT_ID)
        if ctx is None:
            return
        ctx.clear()
        reset_var(CONTEXT_ID)


context_manager = ContextManager()
